from Embed_model import BGE_Embed
from typing import List
from typing_extensions import TypedDict
from pprint import pprint

# LangChain 和 LangGraph 相关导入
from qw_model import qw_model
from langchain_chroma import Chroma
from langchain_core.prompts import ChatPromptTemplate, PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from pydantic import BaseModel, Field
from langgraph.graph import END, StateGraph, START


# 环境设置函数
# def _set_env(key: str):
#     if key not in os.environ:
#         os.environ[key] = getpass.getpass(f"{key}:")
#
#
# _set_env("OPENAI_API_KEY")


# 设置 LangSmith（可选）
# 您可以根据需要配置 LangSmith，以便进行调试和监控
# 详细信息请参考 LangSmith 的官方文档

# -------------------------------------初始化检索器
# 加载知识库
embedding = BGE_Embed()
def initialize_retriever():
    # 添加到向量数据库
    vectorstore = Chroma(
        embedding_function=embedding,
        persist_directory='./chroma_db'
    )

    # 创建检索器
    retriever = vectorstore.as_retriever(search_kwargs={'k': 3})

    return retriever



# 定义图的状态
class GraphState(TypedDict):
    """
    表示图的状态。

    属性：
        question: 用户问题
        generation: LLM生成的回答
        documents: 文档列表
    """
    question: str
    generation: str
    documents: List[str]

# 定义评分器的数据模型和函数

# 1. 检索评分器（Retrieval Grader）
class GradeDocuments(BaseModel):
    """评估检索文档相关性的二元评分。"""
    binary_score: str = Field(
        description="文档是否与问题相关，'yes' 或 'no'"
    )


def initialize_retrieval_grader():
    llm = qw_model()
    structured_llm_grader = llm.with_structured_output(GradeDocuments)

    system = """你是一个评分员，负责评估检索到的
    文档是否与用户的问题相关。
这不需要是严格的测试，目标是过滤掉错误的检索结果。
如果文档包含与用户问题相关的关键词或语义含义，请将其评分为相关。
请给出二元评分“yes”或“no”，以指示文档是否与问题相关。"""
    grade_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", system),
            ("human", "检索到的文档：\n\n {document} \n\n 用户问题：{question}"),
        ]
    )

    retrieval_grader = grade_prompt | structured_llm_grader
    return retrieval_grader


# 2. 幻觉评分器（Hallucination Grader）
class GradeHallucinations(BaseModel):
    """评估回答中是否存在幻觉的二元评分。"""
    binary_score: str = Field(
        description="回答是否基于事实支持，'yes' 或 'no'"
    )


def initialize_hallucination_grader():
    llm = qw_model()
    structured_llm_grader = llm.with_structured_output(GradeHallucinations)

    system = """你是一个评分员，负责评估LLM生成的回答是否基于一组检索到的事实。
请给出二元评分“yes”或“no”。“yes”表示回答是基于这些事实支持的。"""
    hallucination_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", system),
            ("human", "事实集：\n\n {documents} \n\n LLM生成的回答：{generation}"),
        ]
    )

    hallucination_grader = hallucination_prompt | structured_llm_grader
    return hallucination_grader


# 3. 回答评分器（Answer Grader）
class GradeAnswer(BaseModel):
    """评估回答是否解决问题的二元评分。"""
    binary_score: str = Field(
        description="回答是否解决了问题，'yes' 或 'no'"
    )


def initialize_answer_grader():
    llm = qw_model()
    structured_llm_grader = llm.with_structured_output(GradeAnswer)

    system = """你是一个评分员，负责评估一个回答是否解决了用户的问题。
请给出二元评分“yes”或“no”。“yes”表示回答解决了问题。"""
    answer_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", system),
            ("human", "用户问题：\n\n {question} \n\n LLM生成的回答：{generation}"),
        ]
    )

    answer_grader = answer_prompt | structured_llm_grader
    return answer_grader


# 4. 问题重写器（Question Re-writer）
def initialize_question_rewriter():
    llm = qw_model()

    system = """你是一个问题重写器，当用户的问题不够清晰，根据上下文和业务并尝试推理其潜在的语义意图或含义,重新提问相关细节的信息，以优化向量存储的检索效果。
请查看输入的问题。
要牢记公司业务：
1.云剪、实探、真人出镜口播、视频剪辑、扫码直发、电商合作的业务。
2.云剪可以做抖音、快手的平台；实探有小红书、抖音平台；真人口播只有抖音平台，不用询问平台；扫码直发有小红书、抖音，视频剪辑和电商没有平台之分，不用询问平台。
3.抖音平台是按达人等级划分；其他平台是按粉丝数量划分。
4.云剪、实探、扫码直发、真人出镜口播是价格要根据地区是否限制、达人级别（抖音平台）/粉丝数量（其他的平台）去划分的价格。
5.有关视频剪辑的业务价格要根据视频时长、视频复杂程度去定义。
6.电商就粉丝量和用户要做的品类。
"""
    re_write_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", system),
            (
                "human",
                "这是初始问题：\n\n {question} \n 请制定一个改进后的问题。",
            ),
        ]
    )

    question_rewriter = re_write_prompt | llm | StrOutputParser()
    return question_rewriter


# 生成回答（Generate）
def initialize_rag_chain():
    prompt = PromptTemplate(
        template="""你是一个聪明又有亲和力的销售客服，名叫银弓，擅长用轻松、亲切的语气解答用户关于公司业务、价格、产品咨询的问题。你的目标是提供专业、准确的回答，尤其是价格问题，严格按照相应文档回答，风格需要口语化、自然流畅、贴近用户生活，避免书面化和过于正式的表达以及尊称。\n 
                 这是检索到的文档: \n\n {context} \n\n
                 这是用户问题: {question} \n
                 "1. 当用户问题不清楚时，主动询问用户具体需求。"
                 "2. 回答价格问题，严格按照检索的文档给出，不可自行编造。"
                 "3. 当没有用户所说的地区，就按它的省份价格算，当省份也没有那就只能做不限地区的"
                 "4. 保持沟通语气亲切、自然，避免过于正式的措辞，多用口语化的表达和轻松的语气，比如“老板”、“宝子”、“咱们”，禁止使用'您'。"
                 "5. 输出答案应简洁明了，字数在15字左右。" \n
                  需满足以上回答要求。""",
        input_variables=["context", "question"],
    )

    llm = qw_model()

    rag_chain = prompt | llm | StrOutputParser()
    return rag_chain


# 定义节点函数

def retrieve(state, retriever):
    """
    检索文档

    参数：
        state (dict): 当前图的状态
        retriever: 检索器对象

    返回：
        dict: 更新后的状态，包含检索到的文档
    """
    print("---检索---")
    question = state["question"]

    # 检索
    documents = retriever.invoke(question)
    return {"documents": documents, "question": question}


def generate(state, rag_chain):
    """
    生成回答

    参数：
        state (dict): 当前图的状态
        rag_chain: RAG生成链对象

    返回：
        dict: 更新后的状态，包含生成的回答
    """
    print("---生成回答---")
    question = state["question"]
    documents = state["documents"]

    # RAG生成
    generation = rag_chain.invoke({"context": documents, "question": question})
    return {"documents": documents, "question": question, "generation": generation}


def grade_documents(state, retrieval_grader):
    """
    评估检索到的文档是否相关

    参数：
        state (dict): 当前图的状态
        retrieval_grader: 检索评分器对象

    返回：
        dict: 更新后的状态，包含过滤后的相关文档
    """
    print("---检查文档与问题的相关性---")
    question = state["question"]
    documents = state["documents"]

    # 评分每个文档
    filtered_docs = []
    for d in documents:
        score = retrieval_grader.invoke(
            {"question": question, "document": d.page_content}
        )
        grade = score.binary_score
        if grade == "yes":
            print("---评分：文档相关---")
            filtered_docs.append(d)
        else:
            print("---评分：文档不相关---")
            continue
    return {"documents": filtered_docs, "question": question}



def transform_query(state, question_rewriter):
    print("---优化问题---")
    question = state["question"]
    documents = state["documents"]

    # 添加优化次数限制
    optimization_count = state.get("optimization_count", 0)
    if optimization_count > 3:
        return {"documents": documents, "question": question}  # 停止优化

    # 重写问题
    better_question = question_rewriter.invoke({"question": question})
    print(f"优化后的问题：{better_question}")

    # 增加优化次数
    return {"documents": documents, "question": better_question, "optimization_count": optimization_count + 1}


def decide_to_generate(state):
    """
    决定是否生成回答或重新优化问题

    参数：
        state (dict): 当前图的状态

    返回：
        str: 下一个节点的名称
    """
    print("---评估已评分的文档---")
    filtered_documents = state["documents"]

    if not filtered_documents:
        # 所有文档均不相关，重新优化问题
        print("---决策：所有文档与问题不相关，优化问题---")
        return "transform_query"
    else:
        # 有相关文档，生成回答
        print("---决策：生成回答---")
        return "generate"


def grade_generation_v_documents_and_question(state, hallucination_grader, answer_grader):
    """
    评估生成的回答是否基于文档并解决了问题

    参数：
        state (dict): 当前图的状态
        hallucination_grader: 幻觉评分器对象
        answer_grader: 回答评分器对象

    返回：
        str: 下一个节点的名称
    """
    print("---检查幻觉---")
    question = state["question"]
    documents = state["documents"]
    generation = state["generation"]

    score = hallucination_grader.invoke(
        {"documents": documents, "generation": generation}
    )
    grade = score.binary_score

    # 检查是否存在幻觉
    if grade == "yes":
        print("---决策：生成的回答基于文档---")
        # 检查回答是否解决了问题
        print("---评分生成的回答是否解决问题---")
        score = answer_grader.invoke({"question": question, "generation": generation})
        grade = score.binary_score
        if grade == "yes":
            print("---决策：生成的回答解决了问题---")
            return "useful"
        else:
            print("---决策：生成的回答未解决问题---")
            return "not useful"
    else:
        print("---决策：生成的回答未基于文档，重试---")
        return "not supported"


# 构建并编译图
def build_workflow(retrieve_fn, grade_documents_fn, generate_fn, transform_query_fn,
                   decide_to_generate_fn, grade_generation_fn,
                   retriever, rag_chain, retrieval_grader, hallucination_grader, answer_grader, question_rewriter):
    workflow = StateGraph(GraphState)


    # 定义节点
    workflow.add_node("retrieve", lambda state: retrieve_fn(state, retriever))
    workflow.add_node("grade_documents", lambda state: grade_documents_fn(state, retrieval_grader))
    workflow.add_node("generate", lambda state: generate_fn(state, rag_chain))
    workflow.add_node("transform_query", lambda state: transform_query_fn(state, question_rewriter))

    # 构建边
    workflow.add_edge(START, "retrieve")
    workflow.add_edge("retrieve", "grade_documents")
    workflow.add_conditional_edges(
        "grade_documents",
        decide_to_generate_fn,
        {
            "transform_query": "transform_query",
            "generate": "generate",
        },
    )
    workflow.add_edge("transform_query", "retrieve")
    workflow.add_conditional_edges(
        "generate",
        lambda state: grade_generation_fn(state, hallucination_grader, answer_grader),
        {
            "not supported": "generate",
            "useful": END,
            "not useful": "transform_query",
        },
    )

    # 编译图
    app = workflow.compile()
    return app


# 运行图
def run_workflow(app, inputs):
    for output in app.stream(inputs):
        for key, value in output.items():
            # 打印每个节点的状态
            pprint(f"节点 '{key}':")
            # 可选：打印每个节点的详细状态
            # pprint.pprint(value["keys"], indent=2, width=80, depth=None)
        pprint("\n---\n")

    # 打印最终生成的回答
    pprint(value.get("generation", "没有生成回答"))


def main():
    # 初始化组件
    retriever = initialize_retriever()
    retrieval_grader = initialize_retrieval_grader()
    hallucination_grader = initialize_hallucination_grader()
    answer_grader = initialize_answer_grader()
    question_rewriter = initialize_question_rewriter()
    rag_chain = initialize_rag_chain()

    # 构建工作流
    app = build_workflow(
        retrieve_fn=retrieve,
        grade_documents_fn=grade_documents,
        generate_fn=generate,
        transform_query_fn=transform_query,
        decide_to_generate_fn=decide_to_generate,
        grade_generation_fn=grade_generation_v_documents_and_question,
        retriever=retriever,
        rag_chain=rag_chain,
        retrieval_grader=retrieval_grader,
        hallucination_grader=hallucination_grader,
        answer_grader=answer_grader,
        question_rewriter=question_rewriter
    )

    # 运行第一个示例
    print("=== 示例 1 ===")
    inputs1 = {"question": "云剪多少钱"}
    run_workflow(app, inputs1)

    # 运行第二个示例
    print("\n=== 示例 2 ===")
    inputs2 = {"question": "实探不限等级多少钱"}
    run_workflow(app, inputs2)


if __name__ == "__main__":
    main()
